#include "ordinallabelextract.h"

#include <errno.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>

static double
qfunc_data (Data* data)
{
  return qfunc (data->kappa,
                data->logalpha,
                data->n_workers,
                data->n_labels,
                data->tau,
                data->n_images,
                data->gamma,
                data->rho,
                data->loglambda,
                data->pz,
                data->priorz,
                data->ratings,
                data->n_ratings,
                &data->priorkappa.base,
                &data->priorlogalpha.base,
                &data->priortau.base,
                data->priorgamma,
                &data->priorrho.base,
                &data->priorloglambda.base);
}

static void
estep_data (Data* data)
{
  return estep (data->logalpha,
                data->n_labels,
                data->tau,
                data->n_images,
                data->loglambda,
                data->pz,
                data->priorz,
                data->clamp,
                data->ratings,
                data->n_ratings);
}

static double
mstep_data (Data*               data,
            unsigned int        maxiter,
            double              tol)
{
  return mstep (&data->kappa,
                data->logalpha,
                data->n_workers,
                data->n_labels,
                data->tau,
                data->n_images,
                data->gamma,
                &data->rho,
                data->loglambda,
                data->pz,
                data->priorz,
                data->ratings,
                data->n_ratings,
                &data->priorkappa.base,
                &data->priorlogalpha.base,
                &data->priortau.base,
                data->priorgamma,
                &data->priorrho.base,
                &data->priorloglambda.base,
                maxiter,
                tol);
}

static void
em (Data* data)
{
  double q = qfunc_data (data);
  double lastq = q;
  const double epsilon = 1e-5;

  do
    {
      lastq = q;
      estep_data (data);
      mstep_data (data, 25, 1e-3);
      q = qfunc_data (data);
    }
  while (fabs ((q - lastq) / lastq) > epsilon);
}

static void
print_kappa (FILE*       fp,
             const Data* data)
{
  fprintf (fp, "Kappa = %g\n", data->kappa);
}

static void
print_logalpha (FILE*       fp,
                const Data* data)
{
  for (unsigned int i = 0; i < data->n_images; ++i)
    {
      fprintf (fp, "Alpha[%u] = %g\n", i, exp (data->logalpha[i]));
    }
}

static void
print_tau (FILE*        fp,
           const Data*  data)
{
  for (unsigned int i = 0; i < data->n_workers; ++i)
    {
      fprintf (fp, "Tau[%u] =", i);
      for (unsigned int j = 0; j + 1 < data->n_labels; ++j)
        {
          fprintf (fp, " %f", data->tau[i * (data->n_labels - 1) + j]);
        }

      fprintf (fp, "\n");
    }
}

static void
print_gamma (FILE*       fp,
             const Data* data)
{
  fprintf (fp, "Gamma =");
  for (unsigned int i = 0; i + 1 < data->n_labels; ++i)
    {
      fprintf (fp, " %f", data->gamma[i]);
    }
  fprintf (fp, "\n");
}

static void
print_pz (FILE*         fp,
          const Data*   data)
{
  for (unsigned int i = 0; i < data->n_images; ++i)
    {
      fprintf (fp, "P(Z(%u)=(k>0)) =", i);

      for (unsigned int j = 1; j < data->n_labels; ++j)
        {
          fprintf (fp, " %f", data->pz[j + i * data->n_labels]);
        }

      fprintf (fp, "\n");
    }
}

static void
print_rho (FILE*       fp,
             const Data* data)
{
  fprintf (fp, "Kappa = %g\n", data->rho);
}

static void
print_loglambda (FILE*       fp,
                const Data* data)
{
  for (unsigned int i = 0; i < data->n_workers; ++i)
    {
      fprintf (fp, "Lambda[%u] = %g\n", i, exp (data->loglambda[i]));
    }
}

static void
print_parameters (FILE*         fp,
                  const Data*   data)
{
  print_kappa (fp, data);
  print_logalpha (fp, data);
  print_tau (fp, data);
  print_gamma (fp, data);
  print_rho (fp, data);
  print_loglambda (fp, data);
  print_pz (fp, data);
}

static void
usage (void)
{
  fprintf (stderr, 
           "usage: ordinallabelextract [-h] [datafile] [clampfile]\n"
           "\n"
           "  datafile: specify file to read input data from.\n"
           "      if datafile eq '-' use standard input\n"
           "      if not specified, \n"
           "        if standard input is not a tty, use standard input\n"
           "        if standard input is a tty, display this message.\n"
           "\n"
           "  clampfile: specify file to read clamp specification from.\n"
           "      if clampfile eq '-' use standard input\n"
           "        if datafile is also standard input, datafile is read first\n"
           "      if not specified, \n"
           "        use same file as datafile; datafile is read first\n"
           "\n"
           "  -h: display this message.\n"
           "\n"
          );

  exit (1);
}

int 
main (int       argc, 
      char*     argv[])
{
  Data data;

  while (argc > 1)
    {
      if (strcmp (argv[1], "-h") == 0)
        {
          usage ();
        }
      else
        {
          break;
        }
    }

  if (argc > 3 || (argc == 1 && isatty (0)))
    {
      usage ();
    }

  FILE* dfp = (argc < 2 || strcmp (argv[1], "-") == 0) ? stdin
                                                       : fopen (argv[1], "r");

  if (dfp == NULL)
    {
      fprintf (stderr, 
               "can't open %s: %s\n",
               argv[1],
               strerror (errno));
      exit (1);
    }

  if (read_data (dfp, &data) != 0)
    {
      fprintf (stderr, "fatal: read_data failed\n");
      exit (1);
    }

  FILE* cfp = (argc < 3) ? dfp 
                         : strcmp (argv[2], "-") == 0 ? stdin
                                                      : fopen (argv[2], "r");

  if (cfp == NULL)
    {
      fprintf (stderr, 
               "can't open %s: %s\n",
               argv[2],
               strerror (errno));
      exit (1);
    }

  if (read_clamp (cfp, &data) == 0)
    {
      em (&data);
      print_parameters (stdout, &data);
    }

  if (cfp != dfp)
    {
      fclose (cfp);
    }

  fclose (dfp);
  data_destruct (&data);

  return 0;
}
